import base64
from typing import List, Optional
import requests

from config import *
from model import Table


def ocr_img(img: str) -> Optional[List[List[str]]]:
    '''传入base64编码的图片，解析表格'''
    BODY['img'] = img
    r = requests.post(URL, headers=HEADERS, json=BODY)
    
    if r.status_code == 200:
        data = r.json()
        r.close()

        sid = data['sid']
        tables: List[Table] = []
        for table_info in data['prism_tablesInfo']:
            table = Table.from_json(table_info)
            tables.append(table)
    
        for table in tables:
            temp = []
            for _ in range(table.yCellSize):
                temp.append([''] * table.xCellSize)

            for cell in table.cells:
                temp[cell.ysc][cell.xsc] = cell.word

            result = []
            for i in range(table.yCellSize):
                tm = []
                for j in range(table.xCellSize):
                    tm.append(temp[i][j])
                result.append(tm)
            
            return result
    else:
        data = r.json()
        r.close()

        msg = data['error_msg']

        return [[msg]]
    return None


if __name__ == '__main__':
    with open('test.png', 'rb') as f:
        bys = f.read()
        img = base64.b64encode(bys)
        with open('test.txt', 'w', encoding='UTF-8') as f:
            f.write(img.decode(encoding='UTF-8'))
        # try:
        #     ocr_img(img.decode(encoding='UTF-8'))
        # except Exception as e:
        #     pass
